#model.py文件
import torchvision.models as models
import torch.nn as nn

def get_model(num_classes=100):
    # 加载预训练的 ResNet50 模型
    model = models.resnet50(pretrained=True)

    # 修改最后的全连接层以适应 100 类分类任务
    model.fc = nn.Linear(model.fc.in_features, num_classes)

    return model